from bs4 import BeautifulSoup
import requests
from queue import Queue
from time import sleep
import os
from urllib.parse import urljoin
import json
from _thread import start_new_thread
from sys import base_prefix, exit

delay = 0 # second
header={}

running=False
basic_url=""
match_url=""
forbidden_url=[]
vst = {}
vst_img = []

def check_forbidden(a):
	global forbidden_url
	for i in forbidden_url:
		if i in a:
			return True
	return False

def save_to_path(url, img):
	url.replace("?","ASK")
	url.replace("#","POUND")
	url.replace(".","DOT")
	url.replace("=","EQUAL")
	url.replace("&","AND")
	target=url[url.index("//")+2:]
	dirpath=target[:target.rindex("/")]
	# print("Saved to",dirpath,"Target=",target)
	if not os.path.isdir(dirpath):
		os.makedirs(dirpath)
	f=open(target,"wb")
	f.write(img)
	f.close()

def appendToFile(fnm, msg):
        with open(fnm, "a") as f:
                f.write(msg + "\n")
                
allow_suffix = ["jpg", "bmp", "png", "gif", "mp4"]

def checkSuffixAllow(a):
        for i in allow_suffix:
                if i in a:
                        return True
        return False

def save_img(url):
	url.replace("\t","")
	url.replace(" ","")
	if ".jpg?" in url:
		url = url[:url.index("?")]
	print("Img Downloading... url=",url)
	if not checkSuffixAllow(url):
		print("Downloading Failed")
		return
	if check_forbidden(url):
		print("SKIP")
		return
	try:
		sleep(delay)
		resp=requests.get(url,headers=header,stream=False)
		img=resp.content
		resp.close()
		save_to_path(url, img)
	except:
		appendToFile("failure.stt", url)
		print("Downloading Failed")
		return

def save_status():
	global forbidden_url, vst, vst_img
	f=open("mission.stt", "w")
	f.write("{\"basic_url\":\"" + basic_url + "\",\"match_url\":\"" + match_url + "\",\"fbd_url\":"+json.dumps(forbidden_url)+",\"vst\":"+json.dumps(vst)+",\"vst_img\":"+json.dumps(vst_img)+"}")
	f.close()

def load_status():
	global forbidden_url, vst, vst_img, q, basic_url, match_url
	while not q.empty():
		q.get()
	f=open("mission.stt","r")
	tmp=json.loads(f.read())
	basic_url = tmp["basic_url"]
	match_url = tmp["match_url"]
	forbidden_url=tmp["fbd_url"]
	vst=tmp["vst"]
	vst_img=tmp["vst_img"]
	for i in vst:
		if vst[i]==0:
			q.put(i)
def run_once():
	# print("Running...")
	global q
	now=q.get()
	now.replace("\t","")
	now.replace(" ","")
	print("Page Downloading... url=",now, end = "	 ")
	if check_forbidden(now):
		print("SKIP")
		return
	try:
		sleep(delay)
		resp=requests.get(now,headers=header,stream=False)
		r=resp.content.decode("utf-8")
		print("Page downloaded")
		resp.close()
		soup=BeautifulSoup(r,"html.parser")
		for i in soup.find_all('a'):
			if i.has_attr("href"):
				s=i.attrs["href"]
				if s.find("mailto")!=-1:
					continue
				if s.find("http")!=-1:
					if s.find(match_url)!=-1:
						url=s
					else:
						continue
				else:
					url=urljoin(basic_url,s)
				if check_forbidden(url):
					continue
				if url not in vst:
					q.put(url)
					vst.update({url:0})
		for i in soup.find_all("img"):
			if i.has_attr("src"):
				s=i.attrs["src"]
				if s.find("http")!=-1:
					if s.find(match_url)!=-1:
						url=s
					else:
						continue
				else:
					# print("basic_url =", basic_url, "s =", s)
					url=urljoin(basic_url,s)
					# print("url = ", url)
				if url not in vst_img:
					sleep(delay)
					save_img(url)
					vst_img.append(url)
	except:
		print("Downloading Failed")
	vst[now]=1

def run_thread():
	global q,running
	while not q.empty():
		if not running:
			continue
		run_once()
	print("Thread End.")
if __name__ == "__main__":
	q=Queue()

	while True:
		s=input()
		if s=="save":
			save_status()
			print("Status saved")
		if s=="load":
			load_status()
			print("Status loaded")
		if s=="start":
			try:
				start_new_thread(run_thread, ())
				print("Thread started")
			except:
				print("Unable to start thread")
				exit(0)
			running=True
		if s=="resume":
			running=True
			print("Successfully resumed")
		if s=="stop":
			running=False
			print("Successfully stopped")
		if s=="new":
			basic_url=input("Input Basic URL: ")
			match_url=input("Input Match URL: ")
			forbidden_url=[]
			print("Input Forbidden URL:")
			while True:
				s=input()
				if s=="":
					break
				else:
					forbidden_url+=[s]
			vst[basic_url] = 0
			q.put(basic_url)
			print("Successfully created mission")

		if s=="debug":
			print(running)
			print(vst)
			print(vst_img)
